
options (warn = - 1)
library(Matrix)
library(limma)


# The code of the first five supervised gene selection methods refer to 
# "https://github.com/SydneyBioX/scClassify/blob/master/R/featureSelection.R" 


doLimma <- function(exprsMat, cellTypes, exprs_pct = 0.05, pSig=0.001, topN = 50){
  cellTypes <- droplevels(as.factor(cellTypes))
  tt <- list()
  for (i in seq_len(nlevels(cellTypes))) {
    tmp_celltype <- (ifelse(cellTypes == levels(cellTypes)[i], 1, 0))
    design <- stats::model.matrix(~tmp_celltype)
    
    
    meanExprs <- do.call(cbind, lapply(c(0,1), function(i){
      Matrix::rowMeans(exprsMat[, tmp_celltype == i, drop = FALSE])
    }))
    
    meanPct <- do.call(cbind, lapply(c(0,1), function(i){
      Matrix::rowSums(exprsMat[, tmp_celltype == i,
                               drop = FALSE] > 0)/sum(tmp_celltype == i)
    }))
    
    keep <- meanPct[,2] > exprs_pct
    
    y <- methods::new("EList")
    y$E <- exprsMat[keep, ]
    fit <- limma::lmFit(y, design = design)
    fit <- limma::eBayes(fit, trend = TRUE, robust = TRUE)
    tt[[i]] <- limma::topTable(fit, n = Inf, adjust.method = "BH", coef = 2)
    
    
    
    if (!is.null(tt[[i]]$ID)) {
      tt[[i]] <- tt[[i]][!duplicated(tt[[i]]$ID),]
      rownames(tt[[i]]) <- tt[[i]]$ID
    }
    
    tt[[i]]$meanExprs.1 <- meanExprs[rownames(tt[[i]]), 1]
    tt[[i]]$meanExprs.2 <- meanExprs[rownames(tt[[i]]), 2]
    tt[[i]]$meanPct.1 <- meanPct[rownames(tt[[i]]), 1]
    tt[[i]]$meanPct.2 <- meanPct[rownames(tt[[i]]), 2]
  }
  res <- Reduce(union, lapply(tt, function(t) rownames(t[t$logFC > 0 & (t$meanPct.2 - t$meanPct.1) > 0.05 & t$adj.P.Val < pSig,])[seq_len(topN)]))
  
  return(res)
}

doDV <- function(exprsMat, cellTypes, pSig=0.001, topN = 50){
  cellTypes <- droplevels(as.factor(cellTypes))
  tt <- list()
  for (i in seq_len(nlevels(cellTypes))) {
    tmp_celltype <- (ifelse(cellTypes == levels(cellTypes)[i], 1, 0))
    
    meanPct <- do.call(cbind, lapply(c(0,1), function(i){
      Matrix::rowSums(exprsMat[,
                               tmp_celltype == i,
                               drop = FALSE] > 0)/sum(tmp_celltype == i)
    }))
    
    
    posNeg <- (meanPct[,2] - meanPct[,1]) > 0.05
    exprsMat_filt <- exprsMat[posNeg,]
    tt[[i]] <- apply(exprsMat_filt, 1, function(x) {
      df <- data.frame(gene = x, cellTypes = as.factor(tmp_celltype))
      stats::bartlett.test(gene~cellTypes, df)$p.value
    })
    
    tt[[i]] <- stats::p.adjust(tt[[i]], method = "BH")
  }
  tt <- lapply(tt, function(x)sort(x))
  res <- Reduce(union, lapply(tt, function(t) names(t)[seq_len(min(topN, sum(t < pSig)))]))
  return(res)
}

doDD <- function(exprsMat, cellTypes, pSig=0.001, topN = 50){
  cellTypes <- droplevels(as.factor(cellTypes))
  tt <- list()
  for (i in seq_len(nlevels(cellTypes))) {
    tmp_celltype <- ifelse(cellTypes == levels(cellTypes)[i], 1, 0)
    
    
    meanPct <- do.call(cbind, lapply(c(0,1), function(i){
      Matrix::rowSums(exprsMat[,
                               tmp_celltype == i,
                               drop = FALSE] > 0)/sum(tmp_celltype == i)
    }))
    
    posNeg <- (meanPct[,2] - meanPct[,1]) > 0.05
    exprsMat_filt <- exprsMat[posNeg,]
    tt[[i]] <- apply(exprsMat_filt, 1, function(x) {
      x1 <- x[tmp_celltype == 0]
      x2 <- x[tmp_celltype == 1]
      stats::ks.test(x1, x2, alternative = "greater")$p.value
    })
    tt[[i]] <- stats::p.adjust(tt[[i]], method = "BH")
  }
  tt <- lapply(tt, function(x)sort(x))
  res <- Reduce(union, lapply(tt, function(t) names(t)[seq_len(min(topN, sum(t < pSig)))]))
  
  return(res)
}

doDP <- function(exprsMat, cellTypes, threshold = 1, pSig=0.001, topN = 50){
  cellTypes <- droplevels(as.factor(cellTypes))
  tt <- list()
  for (i in seq_len(nlevels(cellTypes))) {
    tmp_celltype <- (ifelse(cellTypes == levels(cellTypes)[i], 1, 0))
    zerosMat <- ifelse(exprsMat > threshold, 1, 0)
    
    tt[[i]] <- apply(zerosMat,1,  function(x){
      tab <- c()
      for (i in c(0,1)) {
        tmp <- factor(x[tmp_celltype == i], levels = c(0, 1))
        tab <- rbind(tab, table(tmp))
      }
      suppressWarnings(stats::chisq.test(tab)$p.value)
    })
    tt[[i]] <- stats::p.adjust(tt[[i]], method = "BH")
  }
  tt = lapply(tt, function(x)sort(x))
  res <- Reduce(union, lapply(tt, function(t) names(t)[seq_len(min(topN, sum(t < pSig)))]))
  return(res)
}

doBI <- function(exprsMat, cellTypes, pSig=0.001, topN = 50){
  cellTypes <- droplevels(as.factor(cellTypes))
  tt <- list()
  for (i in seq_len(nlevels(cellTypes))) {
    tmp_celltype <- (ifelse(cellTypes == levels(cellTypes)[i], 1, 0))
    
    pi <- table(tmp_celltype)/length(tmp_celltype)
    
    agg_mean <- do.call(cbind, lapply(c(0,1), function(i){
      Matrix::rowMeans(exprsMat[, tmp_celltype == i, drop = FALSE])
    }))
    
    agg_sd2 <- do.call(cbind, lapply(c(0,1), function(i){
      apply(exprsMat[, tmp_celltype == i, drop = FALSE], 1, stats::var)
    }))
    
    bi <- abs(agg_mean[,2] - agg_mean[,1])/sqrt(pi[1]*agg_sd2[,1] +
                                                  pi[2]*agg_sd2[,2])
    
    bi <- unlist(bi)
    names(bi) <- rownames(exprsMat)
    bi <- bi[order(bi, decreasing = TRUE)]
    tt[[i]] <- bi
  }
  tt <- lapply(tt, function(x)x)
  res <- Reduce(union, lapply(tt, function(t) names(t)[seq_len(topN)]))
  return(res)
}

library(Seurat)
seurat_2_3 <- function(data, nfeatures = 500){
  seurat_object <- CreateSeuratObject(as.matrix(data))
  seurat_object = FindVariableFeatures(seurat_object, selection.method = "vst", nfeatures = nfeatures, verbose=FALSE)
  vst_genes = seurat_object@assays$RNA@var.features
  seurat_object@assays$RNA@data = log1p(seurat_object@assays$RNA@counts)
  seurat_object = FindVariableFeatures(seurat_object, selection.method = "disp", nfeatures = nfeatures, verbose=FALSE)
  disp_genes = seurat_object@assays$RNA@var.features
  return(list(disp_genes, vst_genes))
}

gene_select <- function(i, epr, label, t_gene){

  gene_set = list()
  a = doLimma(epr[[i]], label[[i]])
  gene_set[[1]] = intersect(a, t_gene)
  cat("Selecting", length(gene_set[[1]]), "genes with DE", "\n")
  a = doDV(epr[[i]], label[[i]])
  gene_set[[2]] = intersect(a, t_gene)
  cat("Selecting", length(gene_set[[2]]), "genes with DV", "\n")
  a = doDD(epr[[i]], label[[i]])
  gene_set[[3]] = intersect(a, t_gene)
  cat("Selecting", length(gene_set[[3]]), "genes with DD", "\n")
  a = doDP(epr[[i]], label[[i]])
  gene_set[[4]] = intersect(a, t_gene)
  cat("Selecting", length(gene_set[[4]]), "genes with DP", "\n")
  a = doBI(epr[[i]], label[[i]])
  gene_set[[5]] = intersect(a, t_gene)
  cat("Selecting", length(gene_set[[5]]), "genes with BI", "\n")
  a = seurat_2_3(epr[[i]])
  gene_set[[6]] = intersect(a[[1]], t_gene)
  cat("Selecting", length(gene_set[[1]]), "genes with disp", "\n")
  gene_set[[7]] = intersect(a[[2]], t_gene)
  cat("Selecting", length(gene_set[[2]]), "genes with vst", "\n")
  
  return(gene_set)
}

library(parallel)
get_result <- function(epr, label, t_gene){
  n = length(epr)
  cat("the number of references is", n, "\n")
  result = mclapply(1:n, gene_select, epr, label, t_gene, mc.cores = 6)
  result = setNames(result, paste0("Ref", 1:n))
  return(result)
}